package com.suemi.missile.client;

import com.alibaba.fastjson.JSONObject;
import com.suemi.missile.common.entity.ESDocument;
import com.suemi.missile.common.utils.Constants;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequestBuilder;
import org.elasticsearch.action.get.*;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.transport.client.PreBuiltTransportClient;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;

public class ESConnection {
  private IndexManager indexer;
  private Pair<String, String> current;
  private int batchSize = 100;

  private final Client client;

  public ESConnection(String host, int port, String cluster) throws UnknownHostException {
    Settings settings = Settings.builder().put("cluster.name", cluster).build();
    client = new PreBuiltTransportClient(settings)
            .addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(host), port));
    indexer = new IndexManager(client);
    specify(Constants.INDEX_NAME, Constants.INDEX_TYPE);
  }

  public ESConnection() throws UnknownHostException {
    this(Constants.ES_HOST, Constants.ES_PORT, Constants.ES_CLUSTER);
  }

  public void specify(String indexName, String indexType) {
    current = new ImmutablePair<String, String>(indexName, indexType);
    if (!indexer.exists(indexName)) {
      indexer.createIndex(indexName, Constants.INDEX_SETTING);
    }

    if (!indexer.exists(indexName, indexType)) {
      indexer.putMapping(indexName, indexType, Constants.INDEX_MAPPING);
    }
  }

  public boolean exists(String id) {
    final GetRequestBuilder req = client.prepareGet(current.getLeft(), current.getRight(), id);
    req.setStoredFields(new String[0]);
    GetResponse resp = req.setRefresh(true).execute().actionGet();
    return resp.isExists();
  }

  public boolean insert(ESDocument document) {
    if (document.documentId() != null && exists(document.documentId())) return false;
    return this.persist(document);
  }

  private List<ESDocument> genFromResp(SearchResponse resp) {
    if(resp == null) return null;
    if(resp.status().getStatus() != 200) return null;
    List<ESDocument> documents = new ArrayList<>();
    for(SearchHit hit : resp.getHits().getHits()) {
      ESDocument document = new ESDocument();
      document.setDocumentId(hit.getId());
      document.fromJsonString(hit.getSourceAsString());
      documents.add(document);
    }
    return documents;
  }

  public List<ESDocument> list(int from, int size) {

   SearchResponse resp = client.prepareSearch(current.getLeft()).setTypes(current.getRight())
           .setQuery(QueryBuilders.matchAllQuery()).setSize(size).setFrom(from).execute().actionGet();

    return genFromResp(resp);
  }

  public List<ESDocument> list(int size) {
    return list(0,size);
  }

  public Long countAll() {
    SearchResponse resp = client.prepareSearch(current.getLeft())
            .setSource(new SearchSourceBuilder().size(0).query(QueryBuilders.matchAllQuery()))
            .execute().actionGet();
    if(resp.status().getStatus() != 200) return null;
    return resp.getHits().getTotalHits();
  }

  public List<ESDocument> findBy(String field, Object value) {
    SearchResponse resp = client.prepareSearch(current.getLeft()).setQuery(QueryBuilders.termQuery(field,value))
            .execute().actionGet();
    return genFromResp(resp);
  }

  public boolean persist(ESDocument document) {
    final IndexRequestBuilder req = client.prepareIndex(current.getLeft(), current.getRight());
    if (document.documentId() != null)
      req.setId(document.documentId());
    JSONObject json = document.toJson();
    json.remove("documentId");
    req.setSource(json.toJSONString(), XContentType.JSON);
    IndexResponse resp = req.execute().actionGet();
    boolean res = (resp.status() == RestStatus.CREATED || resp.status() == RestStatus.OK);
    if (res && document.documentId() == null) document.setDocumentId(resp.getId());
    return res;
  }

  public ESDocument load(String id) {
    final GetRequestBuilder req = client.prepareGet(current.getLeft(), current.getRight(), id);
    GetResponse resp = req.setRefresh(true).execute().actionGet();
    if (!resp.isExists() || resp.isSourceEmpty()) return null;
    String source = resp.getSourceAsString();
    ESDocument doc = new ESDocument();
    doc.setDocumentId(resp.getId());
    doc.fromJsonString(source);
    return doc;
  }

  public List<ESDocument> loadBatch(Collection<String> ids) {
    MultiGetRequestBuilder req = client.prepareMultiGet();
    for (String id : ids)
      req.add(current.getLeft(), current.getRight(), id);
    MultiGetResponse resp = req.get();
    List<ESDocument> res = new ArrayList<>();
    for (MultiGetItemResponse item : resp) {
      if (item.isFailed() || !item.getResponse().isExists() || item.getResponse().isSourceEmpty())
        res.add(null);
      else {
        GetResponse tmp = item.getResponse();
        ESDocument doc = new ESDocument();
        doc.fromJsonString(tmp.getSourceAsString());
        doc.setDocumentId(tmp.getId());
        res.add(doc);
      }
    }
    return res;
  }

  public boolean delete(ESDocument document) {
    return deleteById(document.documentId());
  }

  public void refresh() {
    indexer.refresh(current.getLeft());
  }

  public boolean deleteById(String id) {
    final DeleteRequestBuilder req = client.prepareDelete(current.getLeft(), current.getRight(), id);
    RestStatus status = req.execute().actionGet().status();
    return status == RestStatus.OK || status == RestStatus.NOT_FOUND;
  }

  public List<String> submit(Iterable<ESDocument> documents) {
    int count = 0;
    List<String> res = new ArrayList<>();
    Iterator<ESDocument> iterator = documents.iterator();
    BulkRequestBuilder req = client.prepareBulk();
    while (iterator.hasNext()) {
      ESDocument doc = iterator.next();
      JSONObject json = doc.toJson();
      IndexRequestBuilder builder = client.prepareIndex(current.getLeft(), current.getRight());
      if (doc.documentId() != null)
        builder.setId(doc.documentId());
      json.remove("documentId");
      builder.setSource(json.toJSONString(), XContentType.JSON);
      req.add(builder);
      ++count;

      if (iterator.hasNext() && batchSize > 0 && count < batchSize) continue;

      BulkResponse resp = req.get();
      for (BulkItemResponse item : resp.getItems()) {
        if (item.isFailed()) res.add(null);
        else {
          res.add(item.getId());
          doc.setDocumentId(item.getId());
        }
      }

      req = client.prepareBulk();
      count = 0;
    }
    return res;
  }

  public void destroy() {
    indexer.deleteIndex(current.getLeft());
  }

  public void close() {
    client.close();
  }

  public IndexManager getIndexer() {
    return indexer;
  }
}
